#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include "DPC.h"
#include "tennis_detect.h"
#include "sample_comm_ive.h"
#include "sample_media_ai.h"
#include "vgs_img.h"
#include "misc_util.h"
#include "hi_ive.h"
#include "hi_comm_ive.h"
#include "mpi_ive.h"

using namespace cv;
using namespace std;

struct node {
	int x, y,val;
	node() {}
	node(int n_x, int n_y, int n_val) :x(n_x), y(n_y), val(n_val) {}
	bool operator<(const node& rhs) {
		return val > rhs.val;
	}
};

bool comp(const node & a, const node & b)
{
    return a.val > b.val;
}

static IVE_SRC_IMAGE_S pstSrc;
static IVE_DST_IMAGE_S pstDst;
static IVE_CSC_CTRL_S stCscCtrl;

static HI_VOID IveImageParamCfg(IVE_SRC_IMAGE_S *pstSrc, IVE_DST_IMAGE_S *pstDst,
    VIDEO_FRAME_INFO_S *srcFrame)
{
    pstSrc->enType = IVE_IMAGE_TYPE_YUV420SP;
    // pstSrc->enType = IVE_IMAGE_TYPE_U8C1;
    pstSrc->au64VirAddr[0] = srcFrame->stVFrame.u64VirAddr[0];
    pstSrc->au64VirAddr[1] = srcFrame->stVFrame.u64VirAddr[1];
    pstSrc->au64VirAddr[2] = srcFrame->stVFrame.u64VirAddr[2]; // 2: Image data virtual address

    pstSrc->au64PhyAddr[0] = srcFrame->stVFrame.u64PhyAddr[0];
    pstSrc->au64PhyAddr[1] = srcFrame->stVFrame.u64PhyAddr[1];
    pstSrc->au64PhyAddr[2] = srcFrame->stVFrame.u64PhyAddr[2]; // 2: Image data physical address

    pstSrc->au32Stride[0] = srcFrame->stVFrame.u32Stride[0];
    pstSrc->au32Stride[1] = srcFrame->stVFrame.u32Stride[1];
    pstSrc->au32Stride[2] = srcFrame->stVFrame.u32Stride[2]; // 2: Image data span

    pstSrc->u32Width = srcFrame->stVFrame.u32Width;
    pstSrc->u32Height = srcFrame->stVFrame.u32Height;

    pstDst->enType = IVE_IMAGE_TYPE_U8C3_PACKAGE;
    pstDst->u32Width = pstSrc->u32Width;
    pstDst->u32Height = pstSrc->u32Height;
    pstDst->au32Stride[0] = pstSrc->au32Stride[0];
    pstDst->au32Stride[1] = 0;
    pstDst->au32Stride[2] = 0; // 2: Image data span
}

static HI_VOID IveFrameParamCfg(IVE_SRC_IMAGE_S *pstSrc, IVE_DST_IMAGE_S *pstDst,
    IPC_IMAGE *srcImage, VIDEO_FRAME_INFO_S *srcFrame)
{
    // pstSrc->enType = IVE_IMAGE_TYPE_YUV420SP;
    pstSrc->enType = IVE_IMAGE_TYPE_U8C3_PACKAGE;
    // pstSrc->au64VirAddr[0] = srcImage->u64VirAddr[0];
    // pstSrc->au64VirAddr[1] = srcImage->u64VirAddr[1];
    // pstSrc->au64VirAddr[2] = srcImage->u64VirAddr[2]; // 2: Image data virtual address
    pstSrc->au64VirAddr[0] = srcImage->u64VirAddr;
    // pstSrc->au64VirAddr[1] = 0;
    // pstSrc->au64VirAddr[2] = 0;
    
    pstSrc->au64PhyAddr[0] = srcImage->u64PhyAddr;
    // pstSrc->au64PhyAddr[1] = 0;
    // pstSrc->au64PhyAddr[2] = 0; // 2: Image data physical address

    // pstSrc->au32Stride[0] = srcImage->u32Stride[0];
    // pstSrc->au32Stride[1] = srcImage->u32Stride[1];
    // pstSrc->au32Stride[2] = srcImage->u32Stride[2]; // 2: Image data span
    pstSrc->au32Stride[0] = srcFrame->stVFrame.u32Stride[0];
    pstSrc->au32Stride[1] = srcFrame->stVFrame.u32Stride[1];
    pstSrc->au32Stride[2] = srcFrame->stVFrame.u32Stride[2];
    // pstSrc->au32Stride[0] = 640;
    // pstSrc->au32Stride[1] = 0;
    // pstSrc->au32Stride[2] = 0;
    pstSrc->u32Width = srcImage->u32Width;
    pstSrc->u32Height = srcImage->u32Height;
    // pstSrc->u32Width = srcFrame->stVFrame.u32Width;
    // pstSrc->u32Height = srcFrame->stVFrame.u32Height;

    pstDst->enType = IVE_IMAGE_TYPE_YUV420SP;
    pstDst->u32Width = pstSrc->u32Width;
    pstDst->u32Height = pstSrc->u32Height;
    pstDst->au32Stride[0] = pstSrc->au32Stride[0];
    pstDst->au32Stride[1] = pstSrc->au32Stride[1];
    pstDst->au32Stride[2] = pstSrc->au32Stride[2]; // 2: Image data span
}

static HI_S32 yuvFrame2rgb(VIDEO_FRAME_INFO_S *srcFrame, IPC_IMAGE *dstImage)
{
    IVE_HANDLE hIveHandle;
    HI_S32 s32Ret = 0;
    stCscCtrl.enMode = IVE_CSC_MODE_PIC_BT709_YUV2RGB; // IVE_CSC_MODE_VIDEO_BT601_YUV2RGB
    IveImageParamCfg(&pstSrc, &pstDst, srcFrame);

    s32Ret = HI_MPI_SYS_MmzAlloc_Cached(&pstDst.au64PhyAddr[0], (void **)&pstDst.au64VirAddr[0],
        "User", HI_NULL, pstDst.u32Height*pstDst.au32Stride[0] * 3); // 3: multiple
    if (HI_SUCCESS != s32Ret) {
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        SAMPLE_PRT("HI_MPI_SYS_MmzFree err\n");
        return s32Ret;
    }

    s32Ret = HI_MPI_SYS_MmzFlushCache(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0],
        pstDst.u32Height*pstDst.au32Stride[0] * 3); // 3: multiple
    if (HI_SUCCESS != s32Ret) {
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        SAMPLE_PRT("HI_MPI_SYS_MmzFree_2 err\n");
        return s32Ret;
    }
    // 3: multiple
    memset_s((void *)pstDst.au64VirAddr[0], pstDst.u32Height*pstDst.au32Stride[0] * 3,
        0, pstDst.u32Height*pstDst.au32Stride[0] * 3); // 3: multiple
    HI_BOOL bInstant = HI_TRUE;

    s32Ret = HI_MPI_IVE_CSC(&hIveHandle, &pstSrc, &pstDst, &stCscCtrl, bInstant);
    if (HI_SUCCESS != s32Ret) {
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        return s32Ret;
    }

    if (HI_TRUE == bInstant) {
        HI_BOOL bFinish = HI_TRUE;
        HI_BOOL bBlock = HI_TRUE;
        s32Ret = HI_MPI_IVE_Query(hIveHandle, &bFinish, bBlock);
        while (HI_ERR_IVE_QUERY_TIMEOUT == s32Ret) {
            usleep(100); // 100: usleep time
            s32Ret = HI_MPI_IVE_Query(hIveHandle, &bFinish, bBlock);
        }
    }
    dstImage->u64PhyAddr = pstDst.au64PhyAddr[0];
    dstImage->u64VirAddr = pstDst.au64VirAddr[0];
    dstImage->u32Width = pstDst.u32Width;
    dstImage->u32Height = pstDst.u32Height;

    return HI_SUCCESS;
}

static HI_S32 rgb2yuvFrame(IPC_IMAGE *dstImage, VIDEO_FRAME_INFO_S *srcFrame)
{
    IVE_HANDLE hIveHandle;
    HI_S32 s32Ret = 0;
    stCscCtrl.enMode = IVE_CSC_MODE_PIC_BT709_RGB2YUV; // IVE_CSC_MODE_VIDEO_BT601_RGB2YUV
    IveFrameParamCfg(&pstSrc, &pstDst, dstImage, srcFrame);

    // pstDst.au64PhyAddr[0] = srcFrame->stVFrame.u64PhyAddr[0];
    // pstDst.au64PhyAddr[1] = srcFrame->stVFrame.u64PhyAddr[1];
    // pstDst.au64PhyAddr[2] = srcFrame->stVFrame.u64PhyAddr[2];
    // pstDst.au64VirAddr[0] = srcFrame->stVFrame.u64VirAddr[0];
    // pstDst.au64VirAddr[1] = srcFrame->stVFrame.u64VirAddr[1];
    // pstDst.au64VirAddr[2] = srcFrame->stVFrame.u64VirAddr[2];
    // pstDst.u32Width = srcFrame->stVFrame.u32Width;
    // pstDst.u32Height = srcFrame->stVFrame.u32Height;
    // cout<<"u64PhyAddr ="<<pstDst.au64PhyAddr[0]<<"\n"<<endl;
    // cout<<"u64VirAddr ="<<pstDst.au64VirAddr[0]<<"\n"<<endl;
    // cout<<"au32Stride1 ="<<pstDst.au32Stride[0]<<", au32Stride2 ="<<pstDst.au32Stride[1]<<", au32Stride3 ="<<pstDst.au32Stride[2]<<",Weight ="<<pstDst.u32Width<< ",Height = "<< pstDst.u32Height<<"\n"<<endl;
    // cout<<"au32Stride1 ="<<pstSrc.au32Stride[0]<<", au32Stride2 ="<<pstSrc.au32Stride[1]<<", au32Stride3 ="<<pstSrc.au32Stride[2]<<",Weight ="<<pstSrc.u32Width<< ",Height = "<< pstSrc.u32Height<<"\n"<<endl;
    
    s32Ret = HI_MPI_SYS_MmzAlloc_Cached(&pstDst.au64PhyAddr[0], (void **)&pstDst.au64VirAddr[0],
        "User", HI_NULL, pstDst.u32Height*pstDst.au32Stride[0] * 3); // 3: multiple
    if (HI_SUCCESS != s32Ret) {
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        SAMPLE_PRT("HI_MPI_SYS_MmzFree err\n");
        //  LOG(ERROR)<<"HI_MPI_SYS_MmzAlloc_Cached Failed with 0x"<<hex<<s32Ret<<endl;
        // printf("HI_MPI_SYS_MmzAlloc_Cached Failed with %x",s32Ret);
        // cout<< "HI_MPI_SYS_MmzAlloc_Cached Failed with "<< s32Ret << endl;
        return s32Ret;
    }

    s32Ret = HI_MPI_SYS_MmzFlushCache(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0],
        pstDst.u32Height*pstDst.au32Stride[0] * 3); // 3: multiple
    if (HI_SUCCESS != s32Ret) {
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        return s32Ret;
    }
    // 3: multiple
    memset_s((void *)pstDst.au64VirAddr[0], pstDst.u32Height*pstDst.au32Stride[0] * 3,
        0, pstDst.u32Height*pstDst.au32Stride[0] * 3); // 3: multiple
    HI_BOOL bInstant = HI_TRUE;

    s32Ret = HI_MPI_IVE_CSC(&hIveHandle, &pstSrc, &pstDst, &stCscCtrl, bInstant);
    if (HI_SUCCESS != s32Ret) {
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        return s32Ret;
    }

    if (HI_TRUE == bInstant) {
        HI_BOOL bFinish = HI_TRUE;
        HI_BOOL bBlock = HI_TRUE;
        s32Ret = HI_MPI_IVE_Query(hIveHandle, &bFinish, bBlock);
        while (HI_ERR_IVE_QUERY_TIMEOUT == s32Ret) {
            usleep(100); // 100: usleep time
            s32Ret = HI_MPI_IVE_Query(hIveHandle, &bFinish, bBlock);
        }
    }
    // srcFrame->stVFrame.u64PhyAddr[0] = pstDst.au64PhyAddr[0];
    // srcFrame->stVFrame.u64PhyAddr[1] = pstDst.au64PhyAddr[1];
    // srcFrame->stVFrame.u64PhyAddr[2] = pstDst.au64PhyAddr[2];
    // srcFrame->stVFrame.u64VirAddr[0] = pstDst.au64VirAddr[0];
    // srcFrame->stVFrame.u64VirAddr[1] = pstDst.au64VirAddr[1];
    // srcFrame->stVFrame.u64VirAddr[2] = pstDst.au64VirAddr[2];
    // srcFrame->stVFrame.u32Width = pstDst.u32Width;
    // srcFrame->stVFrame.u32Height = pstDst.u32Height;

    return HI_SUCCESS;
}

static HI_S32 frame2Mat(VIDEO_FRAME_INFO_S *srcFrame, Mat &dstMat)
{
    HI_U32 w = srcFrame->stVFrame.u32Width;
    HI_U32 h = srcFrame->stVFrame.u32Height;
    // cout<<"w = "<<w<<"h = "<<h<<endl;
    int bufLen = w * h * 3;
    HI_U8 *srcRGB = NULL;
    IPC_IMAGE dstImage;
    if (yuvFrame2rgb(srcFrame, &dstImage) != HI_SUCCESS) {
        SAMPLE_PRT("yuvFrame2rgb err\n");
        return HI_FAILURE;
    }
    srcRGB = (HI_U8 *)dstImage.u64VirAddr;
    dstMat.create(h, w, CV_8UC3);
    memcpy_s(dstMat.data, bufLen * sizeof(HI_U8), srcRGB, bufLen * sizeof(HI_U8));
    HI_MPI_SYS_MmzFree(dstImage.u64PhyAddr, (void *)&(dstImage.u64VirAddr));
    return HI_SUCCESS;
}

static HI_S32 Mat2frame(Mat &srcMat, VIDEO_FRAME_INFO_S *desFrame)
{
    HI_U32 w = desFrame->stVFrame.u32Width;
    HI_U32 h = desFrame->stVFrame.u32Height;
    // desFrame->stVFrame.u32Width = 640;
    // desFrame->stVFrame.u32Height = 480;
    // HI_U32 w = 640;
    // HI_U32 h = 480;
    int w2 = srcMat.cols;
    int h2 = srcMat.rows;
    int bufLen = w * h * 3;
    // cout<<"w = "<<w<<", h = "<<h<<", buflen = "<<bufLen<<endl;
    // cout<<"w2 = "<<w2<<", h2 = "<<h2<<", buflen = "<<bufLen<<endl;
    HI_U8 *srcRGB = NULL;
    IPC_IMAGE dstImage;
	srcRGB = (HI_U8 *)dstImage.u64VirAddr;
    dstImage.u32Width = 1920;
    dstImage.u32Height = 1080;
    // srcRGB->u32Width = 1920;
    // srcRGB->u32Height = 1080;
    // cout<<"w3.1 = "<<srcRGB->u32Width<<", h3.1 = "<<srcRGB->u32Height<<endl;
    // cout<< "Image addr = "<<dstImage.u64VirAddr<<"\n"<<endl;
	memcpy_s(srcRGB, bufLen * sizeof(HI_U8), srcMat.data, bufLen * sizeof(HI_U8));
	// cout<<"w3.2 = "<<dstImage.u32Width<<", h3.2 = "<<dstImage.u32Height<<endl;
    // cout<<"hhh"<<endl;
    if (rgb2yuvFrame(&dstImage, desFrame) != HI_SUCCESS) {
        SAMPLE_PRT("rgb2yuvFrame err\n");
        return HI_FAILURE;
    }
    HI_MPI_SYS_MmzFree(dstImage.u64PhyAddr, (void *)&(dstImage.u64VirAddr));
    return HI_SUCCESS;
}

HI_S32 dpc::DPCLoad(uintptr_t* model)
{
    HI_S32 ret = 1;
    *model = 1;
    SAMPLE_PRT("DPC success\n");

    return ret;
}

HI_S32 dpc::DPCUnload(uintptr_t model)
{
    model = 0;

    return HI_SUCCESS;
}

int rows, cols;
//得到暗通道图片
int** getDarkChannel(cv::Mat img,  int blockSize = 3) {
	rows = img.rows;
	cols = img.cols;
	if (img.channels() != 3) {
		fprintf(stderr, "Please use three channel picture input");
		exit(-1);
	}
	//使用指针得到一个二维矩阵
	int** imgGray;
	imgGray = new int* [rows];
	for (int i = 0; i < rows; i++) {
		imgGray[i] = new int[cols];//每一个单位开辟一个新数组
	}
	//求每个像素点的三通道中的最小值，存入刚创建的二维矩阵
	for (int i = 0;i < rows;i++) {
		for (int j = 0;j < cols; j++) {
			int min = 255;
			for (int k = 0;k < 3;k++) {//三次循环求最小值
				if (img.at<Vec3b>(i, j)[k] < min) {
					min = img.at<Vec3b>(i, j)[k];
				}
			}
			imgGray[i][j] = min;
		}
	}//每个像素点的三通道最小值已经求好，得到imgGray矩阵

	//开始求暗通道的第二部步，把刚刚求的二维矩阵进行最小值滤波
	if (blockSize % 2 == 0 || blockSize < 3) {
		fprintf(stderr, "The window for calculating the minimum filter is too small");
		exit(-1);
	}
	int poolSize = (blockSize - 1) / 2;//滤波窗口大小
	int newHeight = rows + blockSize - 1;
	int newWidth = cols + blockSize - 1;
	//使用一个扩充过的矩阵进行滤波
	int** imgMid = new int* [newHeight];
	for (int i = 0;i < newHeight;i++) {
		imgMid[i] = new int[newWidth];
	}
	for (int i = 0;i < newHeight;i++) {
		for (int j = 0;j < newWidth;j++) {
			if (i < rows && j < cols) {
				imgMid[i][j] = imgGray[i][j];//原对应位置像素不变
			}
			else {
				imgMid[i][j] = 255;//扩充的部分用255填充
			}
		}
	}
	//暗通道对应像素值的矩阵
	int** Dark = new int* [rows];
	cv::Mat imgDark(rows, cols, CV_8UC1);//预留出来要显示的暗通道图片，也就是在这个方法中显示一下
	int darkmin = 255;//预先设定最小值
	for(int i = 0;i < rows;i++){
		Dark[i] = new int[cols];
	}
	//最小值滤波
	for (int i = poolSize;i < newHeight - poolSize;i++) {
		for (int j = poolSize;j < newWidth - poolSize;j++) {
			//每次转圈比较之前都将最小值设置为255
			darkmin = 255;
			for (int k = i - poolSize;k < i + poolSize;k++) {
				for (int l = j - poolSize;l < j + poolSize;l++) {
					if (imgMid[k][l] < darkmin) {
						darkmin = imgMid[k][l];
					}
				}
			}
			Dark[i - poolSize][j - poolSize] = darkmin;
			imgDark.at<uchar>(i - poolSize, j - poolSize) = darkmin;
		}
	}
	// cv::imshow(dark_window_name, imgDark);
	return Dark;//返回最小值滤波的结果矩阵
}


int getA(int** Dark, cv::Mat img, float percent = 0.001) {
	int size = rows * cols;
	std::vector<node> vec_node;
	for (int i = 0;i < rows;i++) {
		for (int j = 0;j < cols;j++) {
			node temp;
			temp.x = i;
			temp.y = j;
			temp.val = Dark[i][j];
			vec_node.push_back(temp);
		}
	}
	//给vec内的元素排序，便于寻找前面0.1%的值
	sort(vec_node.begin(), vec_node.end(),comp);
	//1,取前0.1%的暗通道图片的像素位置。
	//2,在原始有雾图像I中寻找对应的具有最高亮度的点的值，作为A值
	//当precent为0的时候就取暗通道排序后的第一个像素点的三通道值的最大值
	int a = 0;
	if (int(percent * size) == 0) {
		for (int i = 0; i < 3; i++) {
			if (img.at<Vec3b>(vec_node[0].x, vec_node[0].y)[i] > a) {
				a = img.at<Vec3b>(vec_node[0].x, vec_node[0].y)[i];
			}
		}
	}
	for (int i = 0; i < int(percent * size); i++) {
		for (int j = 0; j < 3; j++) {
			if (img.at<Vec3b>(vec_node[i].x, vec_node[i].y)[j] > a) {
				a = img.at<Vec3b>(vec_node[i].x, vec_node[i].y)[j];
			}
		}
	}
	return a;
}


//去雾
cv::Mat getFinalImg(cv::Mat img, float omega = 0.95, float t0 = 0.1, int blockSize = 15, float percent = 0.001) {
	int** imgDark = getDarkChannel(img, blockSize = blockSize);//获取暗通道图片
	int a = getA(imgDark, img, percent = percent);//获取透光率
	
	float** transmission = new float* [rows];
	for (int i = 0;i < rows;i++) {
		transmission[i] = new float[cols];
	}
	for (int i = 0; i < rows; i++) {
		for (int j = 0; j < cols; j++) {
			transmission[i][j] = 1 - omega * imgDark[i][j] / a;//公式(12)
			if (transmission[i][j] < 0.1) {
				transmission[i][j] = 0.1;
			}
		}
	}
	cv::Mat dst(img.rows, img.cols, CV_8UC3);
	for (int channel = 0; channel < 3; channel++) {
		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < cols; j++) {
				int temp = (img.at<Vec3b>(i, j)[channel] - a) / transmission[i][j] + a;//公式（22）
				if (temp > 255) {
					temp = 255;
				}
				if (temp < 0) {
					temp = 0;
				}
				dst.at<Vec3b>(i, j)[channel] = temp;
			}
		}
	}
	return dst;
}



HI_S32 dpc::DPCCal(uintptr_t model, VIDEO_FRAME_INFO_S *srcFrm, VIDEO_FRAME_INFO_S *dstFrm)
{
    (void)model;
    int ret = 0;
    RectBox boxs[32] = {0}; // 32: TENNIS_OBJ_MAX
    int j = 0;

    cv::Mat image;
    frame2Mat(srcFrm, image);
    if (image.size == 0) {
        SAMPLE_PRT("image is null\n");
        return HI_FAILURE;
    }
    
    // cv::Mat image = cv::imread("10.jpg");
    // rows = src.rows;
    // cols = src.cols;
    // cout << "size:" << image.size << endl;
    cv::Mat dst = getFinalImg(image);
    cv::resize( dst, dst, cv::Size(1920, 1080),0,0, cv::INTER_LINEAR);
    // cout << "size:" << dst.size << endl;
    // cout << "dest weight:" << dstFrm->stVFrame.u32Width <<endl;
    // cout << "dest height:" << dstFrm->stVFrame.u32Height <<endl;
    if(dst.size != 0)
    {
	    // imwrite("10_res.jpg", dst);
        Mat2frame(dst, dstFrm);
    }
	// imshow("demo_gray", dst);
    // cv::imshow(rgb_window_name, src);
    // cv::imshow(gray_window_name, dst);
    // cv::imwrite("3zxy.jpg", dst);
    // waitKey(0);
}


